#append STAF lib path into sys.path if it is not there.
import sys
STAF_LIB_PATH = '/opt/staf/lib'
if not '/opt/staf/lib' in sys.path:
    sys.path.append(STAF_LIB_PATH)
from PySTAF import STAFResult

class vMachine():
    def __init__(self, logger = None, handle = 'Installation', 
                 machine = {'machine' : 'local', 'vmrun' : '/usr/bin/vmrun',
                       'OS': 'XP', 'image': '/home/zhe/Harddisk_500G/winxp_pro/Windows_XP_Professional.vmx',
                       'guestip': '10.0.1.137', 'timeout': '10m','snapshot': 'Automation'}):
        self.machine = machine
        self.handle = handle
        #Set auto-unmarshall 
        self.handle.setDoUnmarshallResult(1)
        self.logger = logger
        self.status = 'off'
    
    def process_request(self, request_string):
        self.log_info("%s" % (request_string))
        #Submit request on STAF host machine, not the CLIENT!!
        result = self.handle.submit(self.machine['machine'], 'PROCESS', request_string)
        return self.check_result(result, "Error during request, [%s]" % (request_string))
    
    def log_info(self, message):
        self.logger.log(self.logger.Info, message)

    def log_error(self, message):
        self.logger.log(self.logger.Error, message)
        print message

    def log_fail(self, message):
        self.logger.log(self.logger.Fail, message)
        print message
    
    def check_result(self, result_code, error_message):
        if (result_code.rc != STAFResult.Ok):
            if (result_code.rc == 37):
                self.log_error(error_message + ' Request timed out! rc: %d, message: %s' % (result_code.rc, result_code.result))
            else:
                self.log_error(error_message + " rc: %d, message: %s" % (result_code.rc, result_code.result))      
            return False
        
        if result_code.resultObj != None:
            try:
                STDOUT = result_code.resultObj['fileList'][0]['data']
                if STDOUT[:5] == 'Error':
                    #Vmware error
                    self.log_error(error_message + 'Vmware error! rc:%d, message: %s' % (result_code.rc, result_code.resultContext))
                    return -1
            except:
                pass

        return True
        
    def revert(self, snapshot = None):
        """
        
        """
        if not snapshot:
            snapshot = self.machine['snapshot']
        self.logger.log(self.logger.Info, 'Revert to snapshot %s' % snapshot)
        #Set staf to wait for request to finish
        request = 'START COMMAND %s PARMS revertToSnapshot "%s" "%s" \
        WAIT SAMECONSOLE RETURNSTDOUT' % (self.machine['vmrun'], self.machine['image'], self.machine['snapshot'])
        self.logger.log(self.logger.Info, 'Submitting process request: %s' % request)
        result = self.handle.submit(self.machine['machine'], 'PROCESS', request, self.handle.Synchronous)
        if (result.rc != STAFResult.Ok):
            #STAF error
            self.logger.log(self.logger.Error, 
                            'Error starting vmrun process! Result: %s' % (result.resultContext))
            return result.rc
        else:
            #There should be a prettier way to do this
            STDOUT = result.resultObj['fileList'][0]['data']
            if STDOUT[:5] == 'Error':
                #Vmware error
                self.logger.log(self.logger.Error,
                                'Error reverting snapshot! RC:%d, Result: %s' % (result.rc, result.resultContext))
                return -1
            #Revert succeed
            self.logger.log(self.logger.Info,
                            'Revert to snapshot %s succeeded!' % self.machine['snapshot'])
        return result.rc
    
    def start(self):
        self.logger.log(self.logger.Info, 'Start VM [%s]' % self.machine['guestip'])
        #it doesn't work if put 'start' outside the quotation. But now it may not work on Vmware windows.
        request = 'START COMMAND %s PARMS "start %s " WAIT SAMECONSOLE RETURNSTDOUT' % (self.machine['vmrun'], self.machine['image'])
        self.logger.log(self.logger.Info, 'Submitting process request: %s' % request)
        result = self.handle.submit(self.machine['machine'], 'PROCESS', request, self.handle.Synchronous)
        if (result.rc != STAFResult.Ok):
            #STAF error
            self.logger.log(self.logger.Error, 
                            'Error starting vmrun process! Result: %s' % (result.resultContext))
            return result.rc
        else:
            #There should be a prettier way to do this
            STDOUT = result.resultObj['fileList'][0]['data']
            if STDOUT[:5] == 'Error':
                #Vmware error
                self.logger.log(self.logger.Error,
                                'Error starting VM! RC:%d, Result: %s' % (result.rc, result.resultContext))
                return -1
            #Revert succeed
            self.logger.log(self.logger.Info,
                            'Start VM [%s] succeeded!' % self.machine['guestip'])
            self.status = 'on'
            self.updateStatus()
            
        return result.rc
    
    def stop(self):
        parameters = 'stop %s' % self.machine['image']
        if True != self.process_request('START COMMAND "%s" PARMS "%s" WAIT SAMECONSOLE RETURNSTDOUT' % (self.machine['vmrun'], parameters)):
            return False
        self.status = 'off'
        return True
    
    def updateStatus(self):
        #Update VM status. Returns 'online' if OS is up and running.
        self.logger.log(self.logger.Info, 'Start to ping [%s]' % self.machine['guestip'])
        request = 'PING'
        self.logger.log(self.logger.Info, 'Submitting process request: %s' % request)
        result = self.handle.submit(self.machine['guestip'], 'PING', request, self.handle.Synchronous)
        if (result.rc != STAFResult.Ok):
            #STAF error
            self.logger.log(self.logger.Error, 
                            'Error pinging guest machine, OS not ready? Result: %s' % (result.resultContext))
            return 'offline'
        else:
            self.logger.log(self.logger.Info,
                            'Machine [%s] is online.' % self.machine['guestip'])
        #rc=0 when it's online
        self.status = 'online'
        return 'online'